#include "stdafx.h"
//#include "StdAfx.h"
#include "stdio.h"
#include <tchar.h>
#include <atlbase.h>
#include <atlexcept.h>
#include <atlconv.h>
#include "Psapi.h"
//#include "Userenv.h"

#include "HandleProcess.h"
#include "DebugPrintf.h"
#include "ConditionProcessDefine.h"


//////////////////////////////////////////////////////////////////////

typedef HANDLE (WINAPI *CreateToolhelp32SnapshotFunc)(DWORD, DWORD); 
typedef BOOL (WINAPI *Process32FirstFunc)(HANDLE, LPPROCESSENTRY32W); 
typedef BOOL (WINAPI *Process32NextFunc)(HANDLE, LPPROCESSENTRY32W); 


int GetProcessIdFromNameA(char *pszProcessFileName, _Func& FnCompare = NoneCompareFunc())
{
	int nResult = -1;
	int nRetCode = FALSE;

	HMODULE hKernel32Handle = NULL;
	CreateToolhelp32SnapshotFunc fnCreateToolhelp32Snapshot = NULL;
	Process32FirstFunc fnProcess32First = NULL;
	Process32NextFunc fnProcess32Next = NULL;

	HANDLE hProcessSnap = NULL;
	HANDLE hProcess = NULL;
	PROCESSENTRY32W pe32 = {0};

	CONDITION_PROCESS_ERROR_PRINTF0(pszProcessFileName, "Parameter pszProcessFileName is NULL Fail\n");

	hKernel32Handle = LoadLibrary(_T("Kernel32.dll"));
	CONDITION_PROCESS_ERROR_PRINTF1(hKernel32Handle, "LoadLibrary(\"Kernel32.dll\") ErrorCode=%d Fail", GetLastError());

	fnCreateToolhelp32Snapshot = (CreateToolhelp32SnapshotFunc) GetProcAddress(hKernel32Handle, "CreateToolhelp32Snapshot");
	CONDITION_PROCESS_ERROR_PRINTF1(fnCreateToolhelp32Snapshot, "GetProcAddress(\"CreateToolhelp32Snapshot\") ErrorCode=%d Fail\n", GetLastError());

	fnProcess32First = (Process32FirstFunc) GetProcAddress(hKernel32Handle, "Process32FirstW");
	CONDITION_PROCESS_ERROR_PRINTF1(fnProcess32First, "GetProcAddress(\"Process32First\") ErrorCode=%d Fail\n", GetLastError());

	fnProcess32Next = (Process32NextFunc) GetProcAddress(hKernel32Handle, "Process32NextW");
	CONDITION_PROCESS_ERROR_PRINTF1(fnProcess32Next, "GetProcAddress(\"Process32Next\") ErrorCode=%d Fail\n", GetLastError());

	hProcessSnap = (fnCreateToolhelp32Snapshot)(TH32CS_SNAPPROCESS, 0);
	CONDITION_PROCESS_ERROR_PRINTF1(hProcessSnap != INVALID_HANDLE_VALUE, "CreateToolhelp32Snapshot() ErrorCode=%d Fail\n", GetLastError());

	pe32.dwSize = sizeof(PROCESSENTRY32W);
	nRetCode = fnProcess32First(hProcessSnap, &pe32);
	CONDITION_PROCESS_ERROR_PRINTF1(nRetCode, "Process32First() ErrorCode=%d Fail\n", GetLastError());
	
	do
	{
		if(stricmp(CW2A(pe32.szExeFile), pszProcessFileName) == 0)
		{
			if (FnCompare(&pe32) == TRUE)
			{
				nResult = pe32.th32ProcessID;
				goto Exit0;
			}
		}

		nRetCode = (fnProcess32Next)(hProcessSnap, &pe32);
		if(!nRetCode)
		{
			nRetCode = GetLastError();
			CONDITION_PROCESS_ERROR_PRINTF1(nRetCode == ERROR_NO_MORE_FILES, "Process32Next() ErrorCode=%d Fail\n", nRetCode);
			break;
		}
	} while(TRUE);
	
Exit0:

	if(hProcessSnap != INVALID_HANDLE_VALUE)
		CloseHandle(hProcessSnap);
	
	if(hKernel32Handle)
		FreeLibrary(hKernel32Handle);

	return nResult;
}

typedef BOOL (WINAPI *EnumProcessesFunc)(DWORD*, DWORD, DWORD*); 
typedef BOOL (WINAPI *EnumProcessModulesFunc)(HANDLE, HMODULE*, DWORD, LPDWORD); 
typedef DWORD (WINAPI *GetModuleBaseNameFunc)(HANDLE, HMODULE, LPSTR, DWORD); 

int GetProcessIdFromNameNT4(char *pszProcessFileName)
{
	int nResult = -1;
	int nRetCode = FALSE;

	HMODULE hPsapiHandle = NULL;
	EnumProcessesFunc fnEnumProcesses = NULL;
	EnumProcessModulesFunc fnEnumProcessModules = NULL;
	GetModuleBaseNameFunc fnGetModuleBaseName = NULL;

	DWORD dwProcessIdArr[4096], dwNeededSize = 0, dwProcesseNum = 0;
    unsigned int i = 0;
	char szProcessName[MAX_PATH] = "";

	CONDITION_PROCESS_ERROR_PRINTF0(pszProcessFileName, "Parameter pszProcessFileName is NULL Fail\n");

	hPsapiHandle = LoadLibrary(_T("Psapi.dll"));
	CONDITION_PROCESS_ERROR_PRINTF1(hPsapiHandle, "LoadLibrary(\"Psapi.dll\") ErrorCode=%d Fail", GetLastError());

	fnEnumProcesses = (EnumProcessesFunc) GetProcAddress(hPsapiHandle, "EnumProcesses");
	CONDITION_PROCESS_ERROR_PRINTF1(fnEnumProcesses, "GetProcAddress(\"EnumProcesses\") ErrorCode=%d Fail\n", GetLastError());

	fnEnumProcessModules = (EnumProcessModulesFunc) GetProcAddress(hPsapiHandle, "EnumProcessModules");
	CONDITION_PROCESS_ERROR_PRINTF1(fnEnumProcessModules, "GetProcAddress(\"EnumProcessModules\") ErrorCode=%d Fail\n", GetLastError());

	fnGetModuleBaseName = (GetModuleBaseNameFunc) GetProcAddress(hPsapiHandle, "GetModuleBaseNameA");
	CONDITION_PROCESS_ERROR_PRINTF1(fnGetModuleBaseName, "GetProcAddress(\"GetModuleBaseNameA\") ErrorCode=%d Fail\n", GetLastError());

	nRetCode = (fnEnumProcesses)(dwProcessIdArr, sizeof(dwProcessIdArr), &dwNeededSize);
	CONDITION_PROCESS_ERROR_PRINTF1(nRetCode, "EnumProcesses() ErrorCode=%d Fail\n", GetLastError());

    dwProcesseNum = dwNeededSize / sizeof(DWORD);
    for ( i = 0; i < dwProcesseNum; i++ )
	{
		HANDLE hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, dwProcessIdArr[i]);
		if (NULL != hProcess )
		{
			HMODULE hMod;
			DWORD dwNeeded;

			if ( (fnEnumProcessModules)( hProcess, &hMod, sizeof(hMod), &dwNeeded) )
			{
				int nRet = (fnGetModuleBaseName)( hProcess, hMod, szProcessName, sizeof(szProcessName) );
				if(!nRet)
				{
					DEBUG_PRINTF2("GetModuleBaseName() PID=%d ErrorCode=%d Fail\n", dwProcessIdArr[i], GetLastError());
					continue;
				}
			}
			else 
			{
				DEBUG_PRINTF2("EnumProcessModules() PID=%d ErrorCode=%d Fail\n", dwProcessIdArr[i], GetLastError());
				CloseHandle( hProcess );
				continue;
			}
		}
		else
		{
			DEBUG_PRINTF2("OpenProcess(PID=%d) ErrorCode=%d Fail\n", dwProcessIdArr[i], GetLastError());
			continue;
		}

		if(strnicmp(szProcessName, pszProcessFileName, strlen(szProcessName)) == 0)
		{
			nResult = dwProcessIdArr[i];
			CloseHandle(hProcess);
			break;
		}

		CloseHandle( hProcess );
	}

Exit0:

	if(hPsapiHandle)
		FreeLibrary(hPsapiHandle);

	return nResult;
}

int GetProcessIdFromName(char *pszProcessName, _Func& FnCompare)
{
	if(!pszProcessName)
		return -1;

	OSVERSIONINFO osversion = {0};
	osversion.dwOSVersionInfoSize = sizeof(OSVERSIONINFO);

	if (GetVersionEx(&osversion) == false)
	{
		DEBUG_PRINTF1("GetVersionEx() ErrorCode=%d Fail\n", GetLastError());
		return -1;
	}
	else if (osversion.dwPlatformId == VER_PLATFORM_WIN32_NT && 
		osversion.dwMajorVersion <= 4)
	{
		return GetProcessIdFromNameNT4(pszProcessName);
	}
	else
	{
		return GetProcessIdFromNameA(pszProcessName, FnCompare);
	}
}

DWORD GetProcessSessionID(DWORD dwProcessID)
{
	typedef BOOL (__stdcall *ProcessIdToSessionIdT)(DWORD, DWORD*);

	BOOL  bRetCode    = FALSE;
	DWORD dwSessionID = 0;
	
	ProcessIdToSessionIdT pFnProcessIdToSessionId;
	HMODULE hDll = LoadLibrary(_T("Kernel32.dll"));
	CONDITION_PROCESS_ERROR_PRINTF0(hDll != NULL, "LoadLibrary(Kernel32.dll) Fail!");

	pFnProcessIdToSessionId = 
		(ProcessIdToSessionIdT)GetProcAddress(hDll, "ProcessIdToSessionId");
	CONDITION_PROCESS_ERROR_PRINTF0(pFnProcessIdToSessionId != NULL, "GetProcAddress(ProcessIdToSessionId) Fail!");

	bRetCode = pFnProcessIdToSessionId(dwProcessID, &dwSessionID);
	CONDITION_PROCESS_ERROR_PRINTF1(bRetCode != FALSE, "ProcessIdToSessionId Fail : LastError is %d", GetLastError());

Exit0:
	if (hDll != NULL)
	{
		FreeLibrary(hDll);
	}

	return dwSessionID;
}

int TerminateProcessFromNameInSession(char *pszProcessName, DWORD dwSessionID)
{
	CONDITION_ASSERT_RETURN_FALSE_PRINTF0(pszProcessName, "Parameter pszProcessName is NULL Fail\n");
	
	CompareFunc Func(dwSessionID);
	int nPid = GetProcessIdFromName(pszProcessName, Func);
	if(nPid == -1)
	{
		DEBUG_PRINTF0("GetProcessIdFromName() Fail\n");
		return FALSE;
	}

	HANDLE hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, nPid);
	if(!hProcess)
	{
		DEBUG_PRINTF1("OpenProcess() ErrorCode=%d Fail\n", GetLastError());
		return FALSE;
	}

	BOOL bRet = TerminateProcess(hProcess, 0);
	CloseHandle(hProcess);
	return bRet;
}

int TerminateProcessFromName(char *pszProcessName)
{
	CONDITION_ASSERT_RETURN_FALSE_PRINTF0(pszProcessName, "Parameter pszProcessName is NULL Fail\n");

	int nPid = GetProcessIdFromName(pszProcessName);
	if(nPid == -1)
	{
		DEBUG_PRINTF0("GetProcessIdFromName() Fail\n");
		return FALSE;
	}

	HANDLE hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, nPid);
	if(!hProcess)
	{
		DEBUG_PRINTF1("OpenProcess() ErrorCode=%d Fail\n", GetLastError());
		return FALSE;
	}	

	BOOL bRet = TerminateProcess(hProcess, 0);
	CloseHandle(hProcess);
	return bRet;
}

int GetCurrentLogonUserToken(PHANDLE pToken)
{
	int nResult = FALSE;
	int nRetCode = FALSE;

	int nPid = -1;
	HANDLE hProcess = NULL;
	
	CONDITION_PROCESS_ERROR_PRINTF0(pToken, "Parameter pToken is NULL Fail\n");

	nPid = GetProcessIdFromName("explorer.exe");
	CONDITION_PROCESS_ERROR_PRINTF0(nPid >= 0, "GetProcessIdFromName(\"explorer.exe\") Fail\n");

	hProcess = OpenProcess(PROCESS_ALL_ACCESS, FALSE, nPid);
	CONDITION_PROCESS_ERROR_PRINTF1(hProcess, "OpenProcess(\"explorer.exe\") ErrorCode=%d Fail\n", GetLastError());

	nRetCode = OpenProcessToken(
		hProcess, 
		TOKEN_ASSIGN_PRIMARY | TOKEN_DUPLICATE | TOKEN_IMPERSONATE 
		| TOKEN_QUERY | TOKEN_QUERY_SOURCE,//TOKEN_ALL_ACCESS, 
		pToken
	);
	CONDITION_PROCESS_ERROR_PRINTF1(nRetCode, "OpenProcessToken(\"explorer.exe\") ErrorCode=%d Fail\n", GetLastError());

	nResult = TRUE;

Exit0:

	if(hProcess)
		CloseHandle(hProcess);

	return nResult;
}

int LoadUserEnvironmentFunc(PUSER_ENVIRONMENT_FUNC_HANDLE pUserenvHandle)
{
	int nResult = FALSE;
	int nRetCode = FALSE;

	CONDITION_PROCESS_ERROR_PRINTF0(pUserenvHandle, "Parameter pUserenvHandle is NULL Fail\n");

	pUserenvHandle->m_hUserenv = LoadLibrary(_T("Userenv.dll"));
	CONDITION_PROCESS_ERROR_PRINTF1(pUserenvHandle->m_hUserenv, "LoadLibrary(\"Userenv.dll\") ErrorCode=%d Fail\n", GetLastError());

	pUserenvHandle->m_pCreateEnvironmentBlockFunc = (CreateEnvironmentBlockFunc)GetProcAddress(pUserenvHandle->m_hUserenv, "CreateEnvironmentBlock");
	CONDITION_PROCESS_ERROR_PRINTF1(pUserenvHandle->m_pCreateEnvironmentBlockFunc, "GetProcAddress(\"CreateEnvironmentBlock\") ErrorCode=%d Fail\n", GetLastError());

	pUserenvHandle->m_pDestroyEnvironmentBlockFunc = (DestroyEnvironmentBlockFunc)GetProcAddress(pUserenvHandle->m_hUserenv, "DestroyEnvironmentBlock");
	CONDITION_PROCESS_ERROR_PRINTF1(pUserenvHandle->m_pDestroyEnvironmentBlockFunc, "GetProcAddress(\"DestroyEnvironmentBlock\") ErrorCode=%d Fail\n", GetLastError());

	pUserenvHandle->m_pLoadUserProfileFunc = (LoadUserProfileFunc)GetProcAddress(pUserenvHandle->m_hUserenv, "LoadUserProfileA");
	CONDITION_PROCESS_ERROR_PRINTF1(pUserenvHandle->m_pLoadUserProfileFunc, "GetProcAddress(\"LoadUserProfileA\") ErrorCode=%d Fail\n", GetLastError());

	pUserenvHandle->m_pUnloadUserProfileFunc = (UnloadUserProfileFunc)GetProcAddress(pUserenvHandle->m_hUserenv, "UnloadUserProfile");
	CONDITION_PROCESS_ERROR_PRINTF1(pUserenvHandle->m_pUnloadUserProfileFunc, "GetProcAddress(\"UnloadUserProfile\") ErrorCode=%d Fail\n", GetLastError());

	pUserenvHandle->m_hAdvapi32 = LoadLibrary(_T("Advapi32.dll"));
	CONDITION_PROCESS_ERROR_PRINTF1(pUserenvHandle->m_hAdvapi32, "LoadLibrary(\"Advapi32.dll\") ErrorCode=%d Fail\n", GetLastError());

	pUserenvHandle->m_pGetTokenInformationFunc = (GetTokenInformationFunc)GetProcAddress(pUserenvHandle->m_hAdvapi32, "GetTokenInformation");
	CONDITION_PROCESS_ERROR_PRINTF1(pUserenvHandle->m_pGetTokenInformationFunc, "GetProcAddress(\"GetTokenInformation\") ErrorCode=%d Fail\n", GetLastError());

	pUserenvHandle->m_pLookupAccountSidFunc = (LookupAccountSidFunc)GetProcAddress(pUserenvHandle->m_hAdvapi32, "LookupAccountSidA");
	CONDITION_PROCESS_ERROR_PRINTF1(pUserenvHandle->m_pLookupAccountSidFunc, "GetProcAddress(\"LookupAccountSidA\") ErrorCode=%d Fail\n", GetLastError());

//Exit1:
	
	nResult = TRUE;

Exit0:

	if(!nResult)
	{
		if(pUserenvHandle->m_hUserenv)
			FreeLibrary(pUserenvHandle->m_hUserenv);

		if(pUserenvHandle->m_hAdvapi32)
			FreeLibrary(pUserenvHandle->m_hAdvapi32);

		memset(pUserenvHandle, 0, sizeof(USER_ENVIRONMENT_FUNC_HANDLE));
	}
	
	return nResult;
}

int UnloadUserEnvironmentFunc(PUSER_ENVIRONMENT_FUNC_HANDLE pUserenvHandle)
{
	if(pUserenvHandle->m_hUserenv)
	{
		FreeLibrary(pUserenvHandle->m_hUserenv);
	}

	if(pUserenvHandle->m_hAdvapi32)
	{
		FreeLibrary(pUserenvHandle->m_hAdvapi32);
	}

	memset(pUserenvHandle, 0, sizeof(USER_ENVIRONMENT_FUNC_HANDLE));

	return TRUE;
}

int CreateProcessAsCurrentLogonUserFromName(PUSER_ENVIRONMENT_FUNC_HANDLE pUserenvHandle, char *pszProcessFullNameCommand, HANDLE *phProcess)
{	
	int nResult = FALSE;
	int nRetCode = FALSE;

	STARTUPINFOA si;
	PROCESS_INFORMATION stProcessInfo;

	OSVERSIONINFO osvi;

	HANDLE hToken = NULL;
	LPVOID pEnvironment = NULL;
	PROFILEINFO profile = {0};
	char pUserSid[MAX_PATH] = "";
	char szTokenUser[MAX_PATH] = "";
	char szTokenDomain[MAX_PATH] = "";
	DWORD dwUserSize = 0;
	DWORD dwDomainSize = 0;
	SID_NAME_USE sidType;

	if (phProcess != NULL)
	{
		*phProcess = NULL;
	}

	CONDITION_PROCESS_ERROR_PRINTF0(pUserenvHandle, "Parameter pUserenvHandle is NULL Fail\n");
	CONDITION_PROCESS_ERROR_PRINTF0(pUserenvHandle->m_hUserenv, "Parameter pUserenvHandle->m_hUserenv is NULL Fail\n");
	CONDITION_PROCESS_ERROR_PRINTF0(pUserenvHandle->m_hAdvapi32, "Parameter pUserenvHandle->m_hAdvapi32 is NULL Fail\n");
	CONDITION_PROCESS_ERROR_PRINTF0(pszProcessFullNameCommand, "Parameter pszProcessFullNameCommand is NULL Fail\n");

	memset(&osvi, 0, sizeof(OSVERSIONINFO));
	osvi.dwOSVersionInfoSize = sizeof(OSVERSIONINFO);

	nRetCode = GetVersionEx(&osvi);
	CONDITION_PROCESS_ERROR_PRINTF1(nRetCode, "GetVersionEx() ErrorCode=%d Fail\n", GetLastError());
	CONDITION_PROCESS_ERROR_PRINTF0(osvi.dwPlatformId == VER_PLATFORM_WIN32_NT, "Can't Suppot Non-Winnt System (eg : win9x)");
	
	nRetCode = GetCurrentLogonUserToken(&hToken);
	CONDITION_PROCESS_ERROR_PRINTF0(nRetCode && hToken, "GetCurrentLogonUserToken() Fail\n");

	nRetCode = pUserenvHandle->m_pGetTokenInformationFunc(hToken, TokenUser, (void *)pUserSid, sizeof(pUserSid), &dwUserSize);
	CONDITION_PROCESS_ERROR_PRINTF1(nRetCode, "GetTokenInformation(TokenUser) ErrorCode=%d Fail\n", GetLastError());

	dwUserSize = sizeof(szTokenUser);
	dwDomainSize = sizeof(szTokenDomain);
	nRetCode = pUserenvHandle->m_pLookupAccountSidFunc(
		NULL,
		((TOKEN_USER*)pUserSid)->User.Sid,
		szTokenUser,
		&dwUserSize,
		szTokenDomain,
		&dwDomainSize,
		&sidType
	);
	CONDITION_PROCESS_ERROR_PRINTF1(nRetCode, "LookupAccountSid() ErrorCode=%d Fail\n", GetLastError());

	profile.dwSize = sizeof(PROFILEINFO);
	profile.lpUserName = (LPTSTR)szTokenUser;
	nRetCode = pUserenvHandle->m_pLoadUserProfileFunc(hToken, &profile);
	CONDITION_PROCESS_ERROR_PRINTF1(nRetCode, "LoadUserProfile() ErrorCode=%d Fail\n", GetLastError());
	
	nRetCode = pUserenvHandle->m_pCreateEnvironmentBlockFunc(&pEnvironment, hToken, FALSE);
	CONDITION_PROCESS_ERROR_PRINTF1(nRetCode, "CreateEnvironmentBlock() ErrorCode=%d Fail\n", GetLastError());
		
	memset(&si, 0, sizeof(STARTUPINFOA));
	si.cb = sizeof(STARTUPINFOA);
	nRetCode = CreateProcessAsUserA(
		hToken,
		NULL,
		pszProcessFullNameCommand,
		NULL,
		NULL,
		FALSE,
		NORMAL_PRIORITY_CLASS | CREATE_UNICODE_ENVIRONMENT,
		pEnvironment,
		NULL,
		&si,
		&stProcessInfo
	);
	pUserenvHandle->m_pDestroyEnvironmentBlockFunc(pEnvironment);
	CONDITION_PROCESS_ERROR_PRINTF1(nRetCode, "CreateProcessAsUser() ErrorCode=%d Fail\n", GetLastError());

	CloseHandle(stProcessInfo.hThread);


	if (phProcess != NULL)
	{
		*phProcess = stProcessInfo.hProcess;
	}
	else
	{
		CloseHandle(stProcessInfo.hProcess);
	}

	nResult = TRUE;

Exit0:

	if(profile.hProfile)
		pUserenvHandle->m_pUnloadUserProfileFunc(hToken, profile.hProfile);

	if(hToken)
		CloseHandle(hToken);

	return nResult;
}

BOOL CompareFunc::operator () (LPPROCESSENTRY32W pPe32)
{
	if (GetProcessSessionID(pPe32->th32ProcessID) == m_nSessionID)
	{
		return TRUE;
	}

	return FALSE;
}
